use std::{
    future::Future,
    process::exit,
    sync::{
        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
        Arc, LazyLock,
    },
    time::Duration,
};

use arc_swap::ArcSwapOption;
use tokio::runtime::Handle;
use tokio_util::sync::CancellationToken;

#[allow(dead_code)]
pub static WAIT_TASKS: AtomicU64 = AtomicU64::new(0);

#[allow(dead_code)]
pub static PROCESS_EXIT: AtomicBool = AtomicBool::new(false);

#[allow(dead_code)]
pub static SHUTDOWN_TOKEN: LazyLock<CancellationToken> =
    LazyLock::new(CancellationToken::new);

#[derive(Debug)]
pub struct Runtime {
    pub name: String,
    pub affinity_cores: Option<Vec<usize>>,
    pub threads: usize,

    pub runtime: tokio::runtime::Runtime,
}

#[cfg(target_os = "linux")]
#[allow(dead_code)]
fn set_cpu_affinity(field: usize) -> Result<(), String> {
    use nix::{
        sched::{sched_setaffinity, CpuSet},
        unistd::Pid,
    };

    let mut cpu_set = CpuSet::new();
    cpu_set.set(field).unwrap();
    sched_setaffinity(Pid::from_raw(0), &cpu_set).map_err(|e| e.to_string())?;
    Ok(())
}

#[cfg(not(target_os = "linux"))]
#[allow(dead_code)]
fn set_cpu_affinity(field: usize) -> Result<(), String> {
    Ok(())
}

#[allow(dead_code)]
impl Runtime {
    pub fn create_runtime(
        name: String,
        affinity_cores: Option<Vec<usize>>,
        threads: usize,
    ) -> Self {
        let name_clone = name.clone();
        let affinity_cores_clone = affinity_cores.clone();

        let index = AtomicUsize::new(0);

        if let Some(cores) = &affinity_cores {
            if cores.len() != threads {
                panic!(
                    "expect affinity_cores len: {}, but got {}",
                    threads,
                    cores.len()
                );
            }
        }

        let tokio_runtime = tokio::runtime::Builder::new_multi_thread()
            .enable_all()
            .worker_threads(threads)
            .thread_name(name_clone)
            .on_thread_start(move || {
                if let Some(cores) = &affinity_cores_clone {
                    let id = index.fetch_add(1, Ordering::SeqCst);
                    let field = cores[id];
                    set_cpu_affinity(field).unwrap();
                }
            })
            .build()
            .unwrap();

        Runtime {
            name,
            affinity_cores,
            threads,
            runtime: tokio_runtime,
        }
    }

    pub fn get_handle(&self) -> &Handle {
        self.runtime.handle()
    }

    pub fn get_name(&self) -> &str {
        &self.name
    }

    pub fn get_affinity_cores(&self) -> &Option<Vec<usize>> {
        &self.affinity_cores
    }

    pub fn get_thread_num(&self) -> usize {
        self.threads
    }
}

#[allow(dead_code)]
pub struct RuntimeManager {
    pub main_runtime: ArcSwapOption<Runtime>,
    pub worker_runtime: ArcSwapOption<Runtime>,
}

#[allow(dead_code)]
impl RuntimeManager {
    pub fn get_instance() -> &'static RuntimeManager {
        static RUNTIME_MANAGER: LazyLock<RuntimeManager> =
            LazyLock::new(|| RuntimeManager {
                main_runtime: ArcSwapOption::new(None),
                worker_runtime: ArcSwapOption::new(None),
            });

        &RUNTIME_MANAGER
    }

    pub fn init_main_runtime(threads: usize) {
        let runtime =
            Runtime::create_runtime("main".to_string(), None, threads);
        let instance = RuntimeManager::get_instance();
        instance.main_runtime.store(Some(Arc::new(runtime)));
    }

    pub fn init_worker_runtime(
        affinity_cores: Option<Vec<usize>>,
        threads: usize,
    ) {
        let runtime = Runtime::create_runtime(
            "worker".to_string(),
            affinity_cores,
            threads,
        );
        let instance = RuntimeManager::get_instance();
        instance.worker_runtime.store(Some(Arc::new(runtime)));
    }

    pub fn get_main_runtime() -> Arc<Runtime> {
        let instance = RuntimeManager::get_instance();
        let main_runtime = instance.main_runtime.load().clone();
        main_runtime.expect("main_runtime is empty")
    }

    pub fn get_worker_runtime() -> Arc<Runtime> {
        let instance = RuntimeManager::get_instance();
        let worker_runtime = instance.worker_runtime.load().clone();
        worker_runtime.expect("worker_runtime is empty")
    }

    pub fn main_runtime_block_on<F: Future>(future: F) -> F::Output {
        RuntimeManager::get_main_runtime().runtime.block_on(future)
    }

    pub fn shutdown_timeout(wait_ms: u64) {
        PROCESS_EXIT.store(true, Ordering::Relaxed);
        SHUTDOWN_TOKEN.cancel();

        let mut wait_time_ms: u64 = 0;

        loop {
            if WAIT_TASKS.load(Ordering::Relaxed) == 0 {
                break;
            }

            std::thread::sleep(Duration::from_millis(50));
            wait_time_ms += 50;

            // force exit.
            if wait_time_ms > wait_ms {
                break;
            }
        }

        exit(0);
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::atomic::Ordering, time::Duration};
    use tokio::time::sleep;

    use crate::utils::unique_id::UniqueID;

    use super::*;

    #[test]
    fn test_unique_id() {
        let id = UniqueID::new_with_id(0, 1);
        println!("{} {}", id.crc32(), id.crc64());
    }

    // 多个测试并行执行可能使用同一份全局数据结构, 这导致了不确定的行为.
    // 这里runtime和另一个测试用例相互干扰, 非预期的runtime退出导致panic.
    // 请使用 cargo nextest 测试.

    #[test]
    fn test_runtime_manager() {
        RuntimeManager::init_main_runtime(2);
        RuntimeManager::init_worker_runtime(Some(vec![0, 1]), 2);
        RuntimeManager::main_runtime_block_on(async {
            RuntimeManager::get_worker_runtime()
                .get_handle()
                .spawn(async {
                    println!("123");
                    sleep(Duration::from_millis(100)).await;
                    println!("456");
                    sleep(Duration::from_millis(100)).await;
                    println!("789");
                    PROCESS_EXIT.store(true, Ordering::Relaxed);
                });

            let mut exit = false;
            loop {
                if PROCESS_EXIT.load(Ordering::Relaxed) {
                    println!("process done");
                    exit = true;
                }

                sleep(Duration::from_millis(100)).await;
                if exit {
                    break;
                }
            }

            RuntimeManager::shutdown_timeout(0);
        });
    }
}
